package utils

import (
	"net/http"
	"net/url"
	"strings"
	"sync"
)

type SessionManage struct {
	SessionMap map[string]map[string]*http.Cookie
	rw         sync.RWMutex
	HeaderMap  map[string]map[string]string
	rwh        sync.RWMutex
}

func NewSessionManage() *SessionManage {
	return &SessionManage{
		SessionMap: map[string]map[string]*http.Cookie{},
		HeaderMap:  map[string]map[string]string{},
	}
}

// 传入地址获取cookie
func (s *SessionManage) GetCookies(str string) []*http.Cookie {
	host := ""
	if strings.HasPrefix(str, "http") {
		u, err := url.Parse(str)
		if err == nil {
			host = u.Host
		}
	} else {
		host = str
	}
	s.rw.Lock()
	defer s.rw.Unlock()
	cookiesMap := s.SessionMap[host]
	var cookies []*http.Cookie
	for _, cookie := range cookiesMap {
		cookies = append(cookies, cookie)
	}
	return cookies
}

// websocket
func (s *SessionManage) GetCookiesString(str string) string {
	host := ""
	if strings.HasPrefix(str, "http") {
		u, err := url.Parse(str)
		if err == nil {
			host = u.Host
		}
	} else {
		host = str
	}
	s.rw.Lock()
	defer s.rw.Unlock()
	cookiesMap := s.SessionMap[host]
	cookieStr := ""
	i := 0
	for _, cookie := range cookiesMap {
		if i != len(cookiesMap)-1 {
			cookieStr += cookie.Name + "=" + cookie.Value + "; "
		} else {
			cookieStr += cookie.Name + "=" + cookie.Value
		}
		i++
	}
	return cookieStr
}

// 传入地址 更新cookie
func (s *SessionManage) SetCookies(str string, cookies []*http.Cookie) {
	host := ""
	u, err := url.Parse(str)
	if err == nil {
		host = u.Host
	}
	s.rw.Lock()
	defer s.rw.Unlock()
	oldCookieMap := s.SessionMap[host]
	if oldCookieMap == nil {
		oldCookieMap = map[string]*http.Cookie{}
	}
	for _, cookie := range cookies {
		oldCookieMap[cookie.Name] = cookie
	}
	s.SessionMap[host] = oldCookieMap
	return
}
